from typing import Optional, Union

from axiom_boot.conf.manager import Settings
from axiom_boot.di import service, order, conditional_on_setting
from axiom_boot.logging.setup import get_logger
from axiom_boot.scraper.exceptions import DownloadError
from axiom_boot.scraper.interfaces import ScraperMiddleware
from axiom_boot.scraper.models import Target, Response


logger = get_logger(__name__)


@order(100)
@service()
@conditional_on_setting(key="scraper__retry_enabled", expected_value=True)
class RetryMiddleware(ScraperMiddleware):

    def __init__(self, settings: Settings):
        self._settings = settings.scraper

    async def process_request(self, target: Target) -> Optional[Union[Target, Response]]:
        return target

    async def process_response(self, target: Target, response: Response) -> Optional[Union[Target, Response]]:
        if response.status_code in self._settings.retry_on_status:
            return self._retry(target, f"响应状态码 {response.status_code} 触发重试")
        return response

    async def process_exception(self, target: Target, exception: Exception) -> Optional[Union[Target, Response]]:
        if isinstance(exception, DownloadError):
            return self._retry(target, f"下载错误触发重试: {exception}")
        return None

    def _retry(self, target: Target, reason: str) -> Optional[Target]:
        retry_times = target.metadata.get("retry_times", 0)
        max_retries = self._settings.retry_times

        if retry_times < max_retries:
            retry_times += 1
            logger.warning(
                f"正在重试 ({retry_times}/{max_retries}) URL: {target.url}, "
                f"原因: {reason}"
            )
            new_metadata = target.metadata.copy()
            new_metadata["retry_times"] = retry_times
            return target.model_copy(update={"metadata": new_metadata})
        else:
            logger.error(
                f"已达到最大重试次数 ({max_retries}) URL: {target.url}, "
                f"原因: {reason}. 将放弃此请求。"
            )
            return None 